# -*- encoding: utf-8 -*-
"""
@File   : SQLPerformance.py
@Author : @li-xin12345
@Date   : 2022/08/09
"""
import os
import time
import sys


class SQLJmterTest:
    def __init__(self, dir_test):
        self.dir_test = dir_test
        self.dir_result = os.path.join(dir_test, 'test_result')
        self.result_jtl_path = os.path.join(self.dir_result, 'test_result.jtl')
        self.result_output_path = os.path.join(self.dir_result, 'result_output')

    def add_sql_to_jmx(self, url, user, pwd, threads, loops,
                       f_sql, sql_part, test_plan):
        temp_sql_part = os.path.join(self.dir_test, 'temp_sql_part.txt')
        test_plan_file = os.path.join(self.dir_test, 'new_jmx_file.jmx')
        cp_cmd = f'cp {test_plan} {test_plan_file}'
        os.system(cp_cmd)
        with open(f_sql, 'r') as file1:
            contents = file1.read().splitlines()
            sql_count = len(contents)
            for i in range(sql_count):
                sql_statement = contents[i]
                # 修改模板文件
                cp_cmd = f'rm -rf {temp_sql_part};' \
                    f'cp {sql_part} {temp_sql_part}'
                os.system(cp_cmd)
                sed_cmd = f"sed -i 's#{{dbUrl}}#{url}#g' {temp_sql_part};" \
                    f"sed -i 's#{{username}}#{user}#g' {temp_sql_part};" \
                    f"sed -i 's#{{password}}#{pwd}#g' {temp_sql_part};" \
                    f"sed -i 's#{{id}}#{i+1}#g' {temp_sql_part};" \
                    f"sed -i 's#{{threads_num}}#{threads}#g' " \
                    f"{temp_sql_part};" \
                    f"sed -i 's#{{loop_count}}#{loops}#g' " \
                    f"{temp_sql_part};" \
                    f"sed -i 's#{{SQL_statement}}#{sql_statement}#g' " \
                    f"{temp_sql_part};"
                os.system(sed_cmd)
                # 将修改后的模板文件写入执行计划文件
                with open(temp_sql_part, 'r') as file2:
                    lines = file2.readlines()
                    for line in lines:
                        shell_cmd = f"echo -e '{line}' >> {test_plan_file}"
                        os.system(shell_cmd)
                    file2.close()
            file1.close()
        shell_cmd = f"echo  -e '    </hashTree>\n  </hashTree>\n" \
            f"</jmeterTestPlan>' >> {test_plan_file}"
        os.system(shell_cmd)
        # 删除临时文件
        del_cmd = f'rm -rf {temp_sql_part};'
        os.system(del_cmd)
        return test_plan_file

    def get_jmx_result(self, script_path):
        del_cmd = f'rm -rf {self.dir_result}; mkdir {self.dir_result}'
        os.system(del_cmd)
        run_cmd = f'jmeter -n -t {script_path} ' \
            f'-l {self.result_jtl_path} -e -o {self.result_output_path}'
        os.system(run_cmd)

    def read_result(self, sql_f):
        result_file = os.path.join(self.dir_test, 'execute_time_file.txt')
        touch_cmd = f'touch {result_file}'
        os.system(touch_cmd)
        # 读取sql语句
        with open(sql_f, 'r') as file1:
            sql_lists = file1.readlines()
            file1.close()
            sql_count = len(sql_lists)
        # 读取执行结果并将sql语句和执行时间对应输出
        with open(self.result_jtl_path, 'r') as file2:
            contents = file2.read()
            for i in range(1, sql_count+1):
                request = '_'.join(['Request', str(i)])
                request_str = request + ','
                request_count = contents.count(request_str)
                content_lists = contents.splitlines()
                exec_time_sum = 0
                flag = 0
                for content_list in content_lists:
                    if request_str in content_list:
                        if content_list.split(',')[4] == 'OK':
                            exec_time_sum += int(content_list.split(',')[1])
                            flag += 1
                            if flag == request_count & request_count != 0:
                                avg = exec_time_sum / int(request_count)
                                with open(result_file, 'a+') as file3:
                                    file3.write("{0}:{1}\n".format
                                                (sql_lists[i-1].strip('\n'),
                                                 avg))
                                    file3.close()
                        else:
                            with open(result_file, 'a+') as file3:
                                file3.write("{0}:{1}\n".format
                                            (sql_lists[i-1].strip('\n'),
                                             content_list.split(',')[4]))
                                file3.close()
                            break
            file2.close()
        return result_file

    def result_backup(self, threads_number):
        date_res = time.strftime("%Y-%m-%d", time.localtime())
        result_bak = '_'.join(['result', str(threads_number), date_res])
        result_bak_path = os.path.join(os.path.dirname(self.dir_test),
                                       result_bak)
        mv_cmd = f'mv {self.dir_test} {result_bak_path};' \
            f'mv result {result_bak_path}'
        os.system(mv_cmd)


if __name__ == '__main__':
    db_url = sys.argv[1]
    username = sys.argv[2]
    password = sys.argv[3]
    threads_num = sys.argv[4]
    loop_count = sys.argv[5]
    test_result_path = sys.argv[6]
    sql_file = sys.argv[7]
    new_sql_part = sys.argv[8]
    new_test_plan = sys.argv[9]
    print(f"-----检查结果保存路径是否存在，若不存在则创建目录-----")
    check_dir_cmd = f'''if [ -d {test_result_path} ]; 
                        then echo "exists"; 
                        else mkdir -p {test_result_path}; fi'''
    mkdir_result = os.system(check_dir_cmd)
    print(f"-----step1:修改jmx脚本-----")
    test = SQLJmterTest(test_result_path)
    new_jmx_file = test.add_sql_to_jmx(db_url, username, password,
                                       threads_num, loop_count,
                                       sql_file, new_sql_part, new_test_plan)
    print(f'-----step2:执行jmeter脚本{new_jmx_file}-----')
    test.get_jmx_result(new_jmx_file)
    print('-----step3:查看结果并保存-----')
    test.read_result(sql_file)
    test.result_backup(threads_num)
